import os
from mindspore.ops import Primitive
from ..common.utils import get_stack
from ..common import global_manage



class DumpBase:
    def __init__(self) -> None:
        self.real_data = global_manage.get_value("real_data")
        self.dump_path = global_manage.get_value("dump_path")
        self.prefix = "api_checker_"
        self.stack_info = get_stack()
        primitive_attr = dir(Primitive)
        self._op_attr_dict = {
            name for name in primitive_attr if not name.startswith("__")
        }


    def save_info(self, args, kwargs, save_path):
        info_dict = {}
        info_dict["args"] = args
        info_dict["kwargs"] = kwargs
        info_dict["real_data"] = self.real_data
        info_dict["save_path"] = save_path
        info_dict["stack_info"] = self.stack_info
        return info_dict
    
    def set_new_func(self, class_name, index, ori_func):
        setattr(class_name, self.prefix + str(index), ori_func)

    def get_func(self, class_name, index):
        return getattr(class_name, self.prefix + str(index))
    
    def get_save_path(self, name, index):
        return os.path.join(self.dump_path, name, str(index))
    
    def __getattr__(self, name):
        def method_wrapper(*args, **kwargs):
            out = getattr(self.op, name)(*args, **kwargs)
            if out is self.op:
                return self
            return out
        
        if name in self._op_attr_dict:
            if callable(getattr(self.op, name)):
                return method_wrapper
        if hasattr(self.op, name):
            return getattr(self.op, name)
        return super().__getattr__(self, name)
